# Copyright (c) 2008-2009 Greg Harfst http://techguyinmidtown.com/
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish, dis-
# tribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the fol-
# lowing conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
"""
File:   EC2Controller.py
Author: Greg Harfst (greg@techguyinmidtown.com)
Date:   10/10/2008
Notes:  Code to help us control a cluster of EC2 nodes.

We're using the EC2.py file, which I downloaded from Amazon.  I also
borrowed a good amount of ideas from Peter Skomoroch's DataWrangling code.

We provide an object oriented interface, via the EC2Controller object.  
Create an instance of the object, passing all of the key information, etc.
Then you can use that object to start and stop instances, get lists of
instances, etc.
"""

import os
import sys
import time
import EC2

class EC2Controller(object):
    """
    An object that allows you to set up a cluster of nodes on Amazon's
    Elastic Compute Cloud.  We can query what machines are running, start
    and stop them, and configure them for working together as a cluster.
    """
    def __init__(self, aws_access_key_id=None, aws_secret_access_key=None,
                 keyname=None, key_location=None, image_id=None, cluster_size=None,
                 verbose=None):
        """Construct an EC2Controller object.  

        aws_access_key_id     -- string.  access key ID.  if None, we look for the AWS_ACCESS_KEY_ID environment variable.
        aws_secret_access_key -- string.  secret access key.  if None, we look for the AWS_SECRET_ACCESS_KEY environment variable.
        keyname               -- string.  keyname to use.  defaults to 'gsg-keypair'
        key_location          -- string.  path name to the key file.  defaults to ~/.ec2/id_rsa-gsg-keypair
        image_id              -- string.  name of the EC2 image that we will be launching.  defaults to my image id.
        cluster_size          -- int.     number of EC2 machines to start.  defaults to 3.
        verbose               -- boolen.  if True, we print out information to STDOUT while we execute.  default to False.
        """
        self.aws_access_key_id     = aws_access_key_id     if aws_access_key_id     is not None else os.environ['AWS_ACCESS_KEY_ID']
        self.aws_secret_access_key = aws_secret_access_key if aws_secret_access_key is not None else os.environ['AWS_SECRET_ACCESS_KEY']
        self.keyname      = keyname      if keyname      is not None else "gsg-keypair"
        self.key_location = key_location if key_location is not None else os.path.join(os.path.expanduser('~'), '.ec2', 'id_rsa-'+self.keyname)
        self.image_id     = image_id     if image_id     is not None else "ami-ee2eca87"
        self.cluster_size = cluster_size if cluster_size is not None else 3
        self.verbose      = verbose      if verbose      is not None else False

        self.sshconfigName = '/tmp/config'

    def __repr__(self):
        return "EC2 Controller for at most %d nodes of image '%s'" % (self.cluster_size, self.image_id)

    def getInstances(self):
        """Produce a list of EC2Instance objects, containing the instance's running state and DNS names, etc."""
        conn = self._connectToAWS()
        parsed_response = conn.describe_instances().parse()
        instances = []
        for row in (c for c in parsed_response if c[0] == 'INSTANCE'):
            instanceID, imageID, pubIP, privIP, state = row[1:6]
            instances.append(EC2Controller.EC2Instance(instanceID, imageID, pubIP, privIP, state))
        return instances

    def syncFilesToCluster(self, filenames, remoteDir):
        """Use rsync from this machine to one of the cluster machines.  
        Then rsync from that single cluster machine to every other cluster machine.

        filenames  -- [string].  list of filenames that we'll pass as arguments to rsync
        remoteDir  -- string.    remote path that we'll transfer into.
        """
        excludeRules = '--exclude "*~" --exclude ".svn/" --exclude "*.pyc"'
        hosts = self.getRunningPublicHostnames()
        if hosts:
            rootHostname = hosts[0]
            self._runOnRoot('mkdir -p %s' % remoteDir)
            ssh = "ssh -i %s -o 'StrictHostKeyChecking no'" % self.key_location
            files = ' '.join(filenames)
            cmd = 'rsync -e "%s" %s -avz %s root@%s:%s' % (ssh, excludeRules, files, rootHostname, remoteDir)
            if self.verbose: print cmd
            os.system(cmd)
            
            # sync all the files from the root to all other nodes
            remotenames = [os.path.join(remoteDir, os.path.basename(fn)) for fn in filenames]
            remotefiles = ' '.join(remotenames)
            for privateHostname in self.getRunningPrivateHostnames()[1:]:
                # sync particular files, not the whole dir
                cmd = 'rsync -avz %s %s root@%s:%s' % (excludeRules, remotefiles, privateHostname, remoteDir)
                self._runOnRoot(cmd)

    def waitUntilAllAreRunning(self, timeout=300):
        start_time = time.time()
        while (self.getNumRunningHosts() < self.cluster_size):
            time.sleep(15)
            if time.time() - start_time > timeout:
                break
        time.sleep(10) # sometimes it shows running before the ssh daemon is up
        return (self.getNumRunningHosts() >= self.cluster_size)

    def getNumRunningHosts(self):
        """Return the number of running machines."""
        return len(self.getRunningPublicHostnames())

    def getRunningPublicHostnames(self):
        """Get a list of public hostnames for the running instances."""
        return [i.publicName for i in self.getInstances() if i.isRunning()]

    def getRunningPrivateHostnames(self):
        """Get a list of public hostnames for the running instances."""
        return [i.privateName for i in self.getInstances() if i.isRunning()]

    def getRunningInstanceIDs(self):
        """Get a list of instance IDs for the running instances."""
        return [i.instanceID for i in self.getInstances() if i.isRunning()]

    def startCluster(self):
        """Check what's already running.  Then, if we need to start more, call run_instances."""
        if self.verbose: print "startCluster.  Query for the current instances."
        aliveInstances = [i for i in self.getInstances() if i.isRunning() or i.isPending()]
        if len(aliveInstances) < self.cluster_size:
            numToRun = self.cluster_size - len(aliveInstances)
            if self.verbose: print "Starting %d instances..." % (numToRun)
            conn = self._connectToAWS()
            conn.run_instances(imageId=self.image_id, 
                               minCount=max(numToRun/2, 1),
                               maxCount=max(numToRun,   1), 
                               keyName=self.keyname)
            
        else:
            if self.verbose: print "We do not need to start any instances because we already have %d running or pending." % len(aliveInstances)
            
    def stopCluster(self, excludeIDs=None, numToExclude=None):
        """Terminate instances in our cluster.

        excludeIDs    -- [string]. explicitly exclude certain instance IDs from termination.
        numToExclude  -- int.      save the (first) n instances in the instance list.
        """
        if self.verbose: print "stopCluster.  Query for the current instances."
        runningIDs = self.getRunningInstanceIDs()

        if numToExclude and not excludeIDs:
            excludeIDs = runningIDs[0:numToExclude]
        elif not excludeIDs:
            excludeIDs = []

        idsToTerminate = [id for id in runningIDs if id not in excludeIDs]
        if self.verbose: print "Terminate these instances: '%s'" % idsToTerminate
        conn = self._connectToAWS()
        conn.terminate_instances(idsToTerminate)

    def configureFilesOnCluster(self):
        """Copy over the ssh id_rsa all of the remote hsots."""
        self._createSSHConfigFile()
        hosts = self.getRunningPublicHostnames()
        for host in hosts:
            if self.verbose: print "Configure '%s'" % host
            self._scp(host, self.key_location, '.ssh/id_rsa')
            self._scp(host, self.sshconfigName, '.ssh/config')
            self._run(host, 'touch .ssh/authorized_keys')

    def _createSSHConfigFile(self):
        """Write an ssh config file that we'll copy over to ec2."""
        filename = self.sshconfigName
        if self.verbose: print "Write the ssh config file to '%s'." % filename
        outfile  = open(filename, "w")
        print >> outfile, """
Host *
        GSSAPIAuthentication yes
        ForwardX11Trusted yes
        StrictHostKeyChecking no
        KeepAlive yes
"""
        outfile.close()

    def _scp(self, hostname, src, dest):
        """Execute scp from to the hostname.

        hostname  -- string.  external hostname
        src       -- string.  local path to the file we're copying
        dest      -- string.  remote path
        """
        d = {'key':  self.key_location,
             'src':  src,
             'user': 'root',
             'host': hostname,
             'dest': dest}
        cmdline = "scp -C -q -i %(key)s -o 'StrictHostKeyChecking no' '%(src)s' '%(user)s@%(host)s:%(dest)s'" % d
        if self.verbose: print cmdline
        os.system(cmdline)

    def _runOnRoot(self, cmd):
        """Run the command on the root cluster machine.
        cmd  -- string.  the command to run on the remote host.
        """
        hosts = self.getRunningPublicHostnames()
        if hosts:
            self._run(hosts[0], cmd)

    def _runOnAll(self, cmd):
        """Run the command on every host.  We fork to achieve some concurrency.
        cmd  -- string.  the command to run on the remote host.
        """
        hosts = self.getRunningPublicHostnames()
        pids = {}
        for host in hosts:
            pid = os.fork()
            if pid == 0:
                # child process
                self._run(host, cmd)
                os._exit(0)
            else:
                pids[pid] = True
        # wait for the children processes to finish
        while pids.keys():
            pid, _ = os.wait()
            if pid in pids:
                del pids[pid]

    def _run(self, hostname, cmd):
        """Execute a command on a remote machine.

        hostname  -- string.  external hostname
        cmd       -- string.  the command to execute
        """
        d = {'key':  self.key_location,
             'cmd':  cmd,
             'user': 'root',
             'host': hostname}
        cmdline = "ssh -i %(key)s -o 'StrictHostKeyChecking no' '%(user)s@%(host)s' '%(cmd)s'" % d
        if self.verbose: print cmdline
        os.system(cmdline)

    def _connectToAWS(self):
        """Create a connection using the EC2 module to the AWS service."""
        return EC2.AWSAuthConnection(self.aws_access_key_id, self.aws_secret_access_key)


    class EC2Instance(object):
        """Internal class that encapsulates details on an EC2 instance."""
        def __init__(self, instanceID, imageID, publicName, privateName, state):
            self.instanceID   = instanceID
            self.imageID      = imageID
            self.publicName   = publicName
            self.privateName  = privateName
            self.state        = state

        def isRunning(self):
            return self.state == 'running'
        def isPending(self):
            return self.state == 'pending'
        def isTerminated(self):
            return self.state == 'terminated'
        def isShuttingDown(self):
            return self.state == 'shutting-down'

        def __repr__(self):
            return """
EC2Instance( instance ID: %s 
             state:       %s 
             image:       %s
             public  DNS: %s
             private DNS: %s )""" % (self.instanceID, self.state, self.imageID, self.publicName, self.privateName)



def main():
    c = EC2Controller(cluster_size=2, verbose=True)
    print c.getInstances()
    #c.startCluster()
    #c.stopCluster(numToExclude=1)
    #c.stopCluster()

if __name__ == "__main__":
    main()
